package com.webcat.historymemory;

import com.alibaba.fastjson.JSON;
import com.webcat.entity.po.DbConversationChatRecord;
import com.webcat.redis.service.RedisService;
import com.webcat.service.DbConversationChatRecordService;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.stereotype.Component;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.stream.Collectors;

@Component
public class PostgresChatMemory implements ChatMemory {

    @Resource
    private DbConversationChatRecordService dbConversationChatRecordService;

    private final static int maxRecordSize = 100;

    @Resource
    private RedisService redisService;

    @Override
    public void add(String conversationId, List<Message> messages) {
        DbConversationChatRecord dbConversationChatRecord = new DbConversationChatRecord();
        dbConversationChatRecord.setId(generateId());
        dbConversationChatRecord.setConversationId(conversationId);
        messages.forEach(message -> {
            dbConversationChatRecord.setMessageType(message.getMessageType().getValue());
            if (message instanceof UserMessage userMessage) {
                dbConversationChatRecord.setMessage(userMessage.getText());
            }
            else if (message instanceof AssistantMessage assistantMessage) {
                dbConversationChatRecord.setMessage(assistantMessage.getText());
            }
            if (message instanceof ToolResponseMessage toolResponseMessage) {
                List<DbConversationChatRecord.ToolResponse> toolResponseList = toolResponseMessage.getResponses().stream().map(toolResponse -> {
                    DbConversationChatRecord.ToolResponse toolRes = new DbConversationChatRecord.ToolResponse();
                    toolRes.setId(toolResponse.id());
                    toolRes.setToolName(toolResponse.name());
                    toolRes.setToolName(toolResponse.responseData());
                    return toolRes;
                }).toList();
                dbConversationChatRecord.setMessage(JSON.toJSONString(toolResponseList));
            }
        });
        dbConversationChatRecordService.insert(dbConversationChatRecord);
        redisService.addFixedRecord(conversationId, dbConversationChatRecord, maxRecordSize);
        redisService.expire(conversationId, 3600*24);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Object> list = redisService.getAllList(conversationId, 0);
        List<DbConversationChatRecord> dbConversationChatRecordList;
        if (list != null && !list.isEmpty()) {
            dbConversationChatRecordList = list.stream().map(o -> (DbConversationChatRecord)o).collect(Collectors.toList());
        }
        else {
            dbConversationChatRecordList = dbConversationChatRecordService.selectList(conversationId, maxRecordSize);

        }
        if (dbConversationChatRecordList == null || dbConversationChatRecordList.isEmpty()) {
            return Collections.emptyList();
        }
        //dbConversationChatRecordList = rerankList(dbConversationChatRecordList);
        List<Message> messageList = new ArrayList<>();
        dbConversationChatRecordList.forEach(dbConversationChatRecord -> {
            String messageType = dbConversationChatRecord.getMessageType();
            String message = dbConversationChatRecord.getMessage();
            if (message != null && !message.isEmpty()) {
                if (MessageType.USER.getValue().equals(messageType)) {
                    UserMessage userMessage = new UserMessage(message);
                    messageList.add(userMessage);
                }
                else if (MessageType.ASSISTANT.getValue().equals(messageType)) {
                    AssistantMessage assistantMessage = new AssistantMessage(message);
                    messageList.add(assistantMessage);
                }
                if (MessageType.TOOL.getValue().equals(messageType)) {
                    List<DbConversationChatRecord.ToolResponse> toolResponseList = JSON.parseArray(message, DbConversationChatRecord.ToolResponse.class);
                    if (toolResponseList != null && !toolResponseList.isEmpty()) {
                        List<ToolResponseMessage.ToolResponse> toolResponses = toolResponseList.stream().map(toolRes -> new ToolResponseMessage.ToolResponse(toolRes.getId(), toolRes.getToolName(), toolRes.getResponseData())).toList();
                        ToolResponseMessage toolResponseMessage = new ToolResponseMessage(toolResponses);
                        messageList.add(toolResponseMessage);
                    }
                }

            }
        });
        return messageList;
    }

    @Override
    public void clear(String conversationId) {

    }

    private String generateId() {
        return UUID.randomUUID().toString().replaceAll("-", "");
    }

    private List<DbConversationChatRecord> rerankList(List<DbConversationChatRecord> dbConversationChatRecordList) {
        List<DbConversationChatRecord> newDbConversationChatRecordList = new ArrayList<>();
        int len = dbConversationChatRecordList.size();
        for (int i = 0; i < len; i++) {
            if (i < len-1) {
                DbConversationChatRecord preRecord = dbConversationChatRecordList.get(i);
                DbConversationChatRecord nextRecord = dbConversationChatRecordList.get(i+1);
                DbConversationChatRecord tempRecord = preRecord;
                preRecord = nextRecord;
                nextRecord = tempRecord;
                newDbConversationChatRecordList.add(preRecord);
                newDbConversationChatRecordList.add(nextRecord);
            }
        }
        return newDbConversationChatRecordList;
    }
}
